mod constant;
use constant::*;
use base64::{engine::general_purpose, Engine as _};

pub fn encryption(plaintext: String) -> String {
    // 标准密钥和初始化向量
    let key = [0u8; 16];
    let iv = [1u8; 16];

    // PKCS7填充
    let padded_data = pkcs7_pad(plaintext.as_bytes());

    // 创建ZUC密码实例
    let mut zuc_cipher = ZucCipher::new(key, iv);

    // 分块加密
    let encrypted: Vec<u8> = padded_data
        .chunks(4)
        .flat_map(|chunk| {
            // 将chunk转换为u32，大端序
            let mut block = [0u8; 4];
            block[..chunk.len()].copy_from_slice(chunk);
            let block_u32 = u32::from_be_bytes(block);

            // 生成密钥流并异或加密
            let key_stream = zuc_cipher.next_keystream();
            (block_u32 ^ key_stream).to_be_bytes()
        })
        .collect();

    // Base64编码
    general_purpose::STANDARD.encode(encrypted)
}

/// PKCS7填充实现
fn pkcs7_pad(data: &[u8]) -> Vec<u8> {
    let block_size = 4;
    let pad_value = block_size - (data.len() % block_size);
    let mut padded = data.to_vec();
    padded.extend(vec![pad_value as u8; pad_value]);
    padded
}

struct ZucCipher {
    /// 线性反馈移位寄存器状态
    lfsr: [u32; 17],
    /// 内部寄存器R1
    r1: u32,
    /// 内部寄存器R2
    r2: u32,
}

impl ZucCipher {
    /// 创建新的ZUC密码实例
    fn new(key: [u8; 16], iv: [u8; 16]) -> Self {
        let mut cipher = Self::key_schedule(&key, &iv);
        cipher.initialize();
        cipher
    }

    /// 密钥调度：初始化LFSR状态
    fn key_schedule(key: &[u8], iv: &[u8]) -> Self {
        let mut lfsr = [0u32; 17];
        Self::generate_lfsr_state(&mut lfsr, key, iv);

        Self {
            lfsr,
            r1: 0,
            r2: 0,
        }
    }

    /// 生成LFSR初始状态
    fn generate_lfsr_state(lfsr: &mut [u32; 17], key: &[u8], iv: &[u8]) {
        lfsr.iter_mut()
            .zip(key.iter().zip(KD.iter().zip(iv.iter())))
            .for_each(|(s, (&k, (&kd, &iv)))| {
                *s = ((k as u32) << 23) | ((kd & 0x7fff) << 8) | (iv as u32);
            });
    }

    /// 密码初始化
    fn initialize(&mut self) {
        // 32轮初始化
        for _ in 0..32 {
            let (x0, x1, x2, _x3) = self.bit_reconstruction();
            let w = self.non_linear_transform(x0, x1, x2);
            self.lfsr_initial_mode(w >> 1);
        }

        // 最后一轮处理
        let (x0, x1, x2, _x3) = self.bit_reconstruction();
        self.non_linear_transform(x0, x1, x2);
        self.lfsr_work_mode();
    }

    /// 生成下一个密钥流
    fn next_keystream(&mut self) -> u32 {
        let (x0, x1, x2, x3) = self.bit_reconstruction();
        let keystream = self.non_linear_transform(x0, x1, x2) ^ x3;
        self.lfsr_work_mode();
        keystream
    }

    /// 非线性变换函数
    fn non_linear_transform(&mut self, x0: u32, x1: u32, x2: u32) -> u32 {
        // 计算中间值w
        let w = (x0 ^ self.r1).wrapping_add(self.r2);
        let w1 = self.r1.wrapping_add(x1);
        let w2 = self.r2 ^ x2;

        // S盒和非线性变换
        self.r1 = Self::s_box(Self::l1_transform((w1 << 16) | (w2 >> 16)));
        self.r2 = Self::s_box(Self::l2_transform((w2 << 16) | (w1 >> 16)));

        w
    }

    /// 位重构
    fn bit_reconstruction(&self) -> (u32, u32, u32, u32) {
        const HIGH_MASK: u32 = 0x7fff_8000;
        const LOW_MASK: u32 = 0xffff;
        (
            ((self.lfsr[15] & HIGH_MASK) << 1) | (self.lfsr[14] & LOW_MASK),
            (self.lfsr[9] >> 15) | (self.lfsr[11] << 16),
            (self.lfsr[5] >> 15) | (self.lfsr[7] << 16),
            (self.lfsr[0] >> 15) | (self.lfsr[2] << 16),
        )
    }

    /// 初始模式下的LFSR变换
    fn lfsr_initial_mode(&mut self, u: u32) {
        let lfsr = &mut self.lfsr;

        // 复杂的模31加法和位移运算
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[0], 8), lfsr[0]);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[4], 20), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[10], 21), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[13], 17), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[15], 15), v);

        // 特殊处理零值情况
        let s16 = Self::add_mod31(u, v);
        lfsr[lfsr.len() - 1] = if s16 == 0 {
            LFSR_MASK
        } else {
            s16
        };

        // LFSR移位
        for i in 0..(lfsr.len() - 1) {
            lfsr[i] = lfsr[i + 1];
        }
    }

    /// 工作模式下的LFSR变换
    fn lfsr_work_mode(&mut self) {
        let lfsr = &mut self.lfsr;

        // 类似初始模式，但不添加额外输入
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[0], 8), lfsr[0]);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[4], 20), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[10], 21), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[13], 17), v);
        let v = Self::add_mod31(Self::rotate_left_31(lfsr[15], 15), v);

        // 特殊处理零值情况
        lfsr[lfsr.len() - 1] = if v == 0 {
            LFSR_MASK
        } else {
            v
        };

        // LFSR移位
        for i in 0..(lfsr.len() - 1) {
            lfsr[i] = lfsr[i + 1];
        }
    }

    /// 31位左循环移位
    fn rotate_left_31(a: u32, k: usize) -> u32 {
        ((a << k) | (a >> (31 - k))) & LFSR_MASK
    }

    /// 模31加法
    fn add_mod31(a: u32, b: u32) -> u32 {
        let x = a.wrapping_add(b);
        (x & LFSR_MASK) + (x >> 31)
    }

    /// L1非线性变换
    fn l1_transform(x: u32) -> u32 {
        x ^
            x.rotate_left(2) ^
            x.rotate_left(10) ^
            x.rotate_left(18) ^
            x.rotate_left(24)
    }

    /// L2非线性变换
    fn l2_transform(x: u32) -> u32 {
        x ^
            x.rotate_left(8) ^
            x.rotate_left(14) ^
            x.rotate_left(22) ^
            x.rotate_left(30)
    }

    /// S盒变换
    fn s_box(x: u32) -> u32 {
        let idx = x.to_be_bytes();
        let (y0, y1, y2, y3) = (
            S0[idx[0] as usize],
            S1[idx[1] as usize],
            S0[idx[2] as usize],
            S1[idx[3] as usize]
        );
        u32::from_be_bytes([y0, y1, y2, y3])
    }
}

impl Iterator for ZucCipher {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        Some(self.next_keystream())
    }
}